package com.squirrel.michale.answer;


import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;

/**
 * TODO function description
 *
 * @author m30026548
 * @since 2025-04-08
 */
public class Router {

    private record Packet(int source, int destination, int timestamp) {
    }

    private record Pair(List<Integer> timestamps, int head) {
    }

    private final int memoryLimit;
    private final Queue<Packet> packetQ = new ArrayDeque<>(); // Packet 队列
    private final Set<Packet> packetSet = new HashSet<>(); // Packet 集合
    private final Map<Integer, Pair> destToTimestamps = new HashMap<>(); // destination -> [[timestamp], head]

    public Router(int memoryLimit) {
        this.memoryLimit = memoryLimit;
    }

    public boolean addPacket(int source, int destination, int timestamp) {
        Packet packet = new Packet(source, destination, timestamp);
        if (!packetSet.add(packet)) {
            return false;
        }
        if (packetQ.size() == memoryLimit) { // 太多了
            forwardPacket();
        }
        packetQ.add(packet); // 入队
        destToTimestamps.computeIfAbsent(destination, k -> new Pair(new ArrayList<>(), 0)).timestamps.add(timestamp);
        return true;
    }

    public int[] forwardPacket() {
        if (packetQ.isEmpty()) {
            return new int[]{};
        }
        Packet packet = packetQ.poll(); // 出队
        packetSet.remove(packet);
        destToTimestamps.compute(packet.destination, (k, p) -> new Pair(p.timestamps, p.head + 1)); // 队首下标加一，模拟出队
        return new int[]{packet.source, packet.destination, packet.timestamp};
    }

    public int getCount(int destination, int startTime, int endTime) {
        Pair p = destToTimestamps.get(destination);
        if (p == null) {
            return 0;
        }
        int left = lowerBound(p.timestamps, startTime, p.head - 1);
        int right = lowerBound(p.timestamps, endTime + 1, p.head - 1);
        return right - left;
    }

    private int lowerBound(List<Integer> nums, int target, int left) {
        int right = nums.size();
        while (left + 1 < right) {
            int mid = (left + right) >>> 1;
            if (nums.get(mid) >= target) {
                right = mid;
            } else {
                left = mid;
            }
        }
        return right;
    }


}
